# -*- coding: utf-8 -*-

from numpy import ndarray, vstack
from gensim.models import Doc2Vec
from sklearn.linear_model import LogisticRegression
from doc_to_vector_model.tagged_sentence import TaggedSentence
from sklearn.externals import joblib

from config import pos_path, neg_path, data_path


class Sentiment:
    def __init__(self):
        self.profiler = Doc2Vec(min_count=1, window=10, size=100, sample=1e-4, negative=5, workers=7)
        self.classifier = LogisticRegression()
        LogisticRegression(C=1.0, class_weight=None, dual=False, fit_intercept=True,
                           intercept_scaling=1, penalty='l2', random_state=None, tol=0.0001)
        self.sources = {
            'pos': pos_path,
            'neg': neg_path
        }
        self.sentence_source = TaggedSentence(self.sources)
        self.pos_source = TaggedSentence({'pos': pos_path})
        self.neg_source = TaggedSentence({'neg': neg_path})
        self.sentence_vec = ndarray((0, 100))  # size=100
        self.label_vec = []

    def train(self):
        print('train profiler')
        self.profiler.build_vocab(self.sentence_source.to_list())
        for epoch in range(10):
            self.profiler.train(self.sentence_source.sentences_perm())
        print('build train vec')
        # pos
        for sentence, prefix in self.pos_source:
            vector = self.profiler.docvecs[prefix]
            self.sentence_vec = vstack((self.sentence_vec, vector))
            self.label_vec.append(1)  # pos
        # neg
        for sentence, prefix in self.neg_source:
            vector = self.profiler.docvecs[prefix]
            self.sentence_vec = vstack((self.sentence_vec, vector))
            self.label_vec.append(-1)  # pos
        print('train classifier')
        self.classifier.fit(self.sentence_vec, self.label_vec)

    def dev(self):
        pass

    def predict(self, sentence):
        pass

    def test(self):
        i = 0
        for sentence, prefix in self.sentence_source:
            vector = self.profiler.docvecs[prefix]
            sentiment = self.classifier.predict(vector)
            print(sentence[0], sentiment[0])
            i += 1
            if i > 10:
                break

    def save(self, **kwargs):
        # save word2vec
        filename = kwargs.get('doc2vec')
        if filename:
            filename = data_path(filename)
        else:
            filename = data_path('doc2vec.model')
        self.profiler.save(filename)
        # save bayes
        filename = kwargs.get('logistic')
        if filename:
            filename = data_path(filename)
        else:
            filename = data_path('logistic')
        joblib.dump(self.classifier, filename)

    def load(self, **kwargs):
        # load word2vec
        filename = kwargs.get('doc2vec')
        if filename:
            filename = data_path(filename)
        else:
            filename = data_path('doc2vec.model')
        self.profiler = Doc2Vec.load(filename)
        # load bayes
        filename = kwargs.get('logistic')
        if filename:
            filename = data_path(filename)
        else:
            filename = data_path('logistic')
        self.classifier = joblib.load(filename)

if __name__ == '__main__':
    s = Sentiment()
    s.train()
    s.save()
    # s.load()
    # s.test()
